Skip to content

Conversation

@lingtonglu
Copy link
Contributor

Database & Models

  • Migration: alembic/versions/create_wallets_table.py - Creates the wallets table with all required fields. Add extra column billable to provider_key
  • Model: app/models/wallet.py - Wallet model with relationship to User
  • User relationship: Updated to include wallet relationship
  • Provider Key: add extra billable flag

Core Services

  • WalletService: app/services/wallet_service.py - Implements all core functions:
    • ensure_wallet() - Creates wallet if doesn't exist
    • precheck() - Checks balance > 0 and not blocked
    • adjust() - Handles deposits/deductions with optimistic locking
    • set_blocked() - Block/unblock accounts
    • get() - Retrieve wallet details

API Integration

  • Proxy routes: Added wallet precheck to all API endpoints (/chat/completions, /completions, /images/generations, /images/edits, /embeddings)
  • Usage tracker: Integrated automatic balance deduction after API usage calculation
  • Wallet routes: Basic endpoints to check balance and status at /v1/wallet/

Webhook Infrastructure

  • Stripe webhooks: Added placeholder webhook handlers for:
    • payment_intent.succeeded - For crediting wallet balance
    • payment_intent.payment_failed - For handling failed payments
    • invoice.payment_failed - For blocking accounts on failed subscriptions

Request Flow

  1. Precheck: For API request, if the provider is billable, checks balance > 0 AND NOT blocked
  2. Processing: Request proceeds if precheck passes
  3. Deduction: After completion, usage cost is deducted from wallet balance if provider is billable

String, nullable=True
) # Allow custom base URLs for some providers
model_mapping = Column(JSON, nullable=True) # JSON dict for model name mappings
billable = Column(Boolean, nullable=False, default=False)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shall we add a "billable" annotation at the model level instead of the provider level?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's an open discussion. My current implementation is based on the assumption that the billing object is at "provider" level. You could have "free model" by setting the price to be zero.

@lingtonglu lingtonglu merged commit 25cc95b into TensorBlock:main Sep 2, 2025
3 checks passed
@lingtonglu lingtonglu deleted the balance_system branch September 2, 2025 04:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants